{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e344c84",
   "metadata": {},
   "source": [
    "## Self Supervised Pretraining for Tabular Data\n",
    "\n",
    "We have implemented two Self Supervised Pre-training routines that allow the user to pre-train *all* tabular models in the library with the exception of the TabPerceiver (which is a special monster).\n",
    "\n",
    "The two routines implemented are illustrated in the figures below. The 1st is from [TabNet: Attentive Interpretable Tabular Learning](https://arxiv.org/abs/1908.07442) and is designed for models that do not use transformer-based architectures, while the second is from [SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training](https://arxiv.org/abs/2106.01342), and is designed for models that use transformer-based architectures.\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img width=\"800\" src=\"figures/self_supervised_tabnet.png\">\n",
    "</p>\n",
    "\n",
    "Fig 1. Figure 2 in their paper. I have included de original caption in case is useful, althought the Figure itself is pretty self explanatory\n",
    "\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img width=\"600\" src=\"figures/self_supervised_saint.png\">\n",
    "</p>\n",
    "\n",
    "Fig 2. Figure 1 in their paper. Here the caption is necessary 😏\n",
    "\n",
    "\n",
    "It is beyond the scope of this notebook to explain in detail those implementations. Therefore, we strongly recommend the user to go and read the papers if this functionality is of interest to her/him.\n",
    "\n",
    "One thing is worth noticing however. As seen in Fig 1(the TabNet paper's Fig 2) the masking of the input features happens in the feature space. However, the implementation in this library is inspired by that at the [dreamquark-ai](https://github.com/dreamquark-ai/tabnet) repo, which is in itself inspired by the original implementation (by the way, at this point I will write it once again. All TabNet related things in this library are inspired when not directly based in the code in that repo, therefore, ALL CREDIT TO THE GUYS AT dreamquark-ai).\n",
    "\n",
    "In that implementation the masking happens in the embedding space, and currently does not mask the entire embedding (i.e. categorical feature). We decided to release as it is in this version and we will implement the exact same process described in the paper in future releases. \n",
    "\n",
    "Having said all of the above let's see how to use self supervision for tabular data with `pytorch-widedeep`. We will concentrate in this notebook on the 1st of the two approaches (the 'TabNet approach'). For details on the second approach please see `16_Self_Supervised_Pretraning_pt2`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6f7fdb4",
   "metadata": {},
   "source": [
    "### Self Supervision for non-transformer-based models..\n",
    "\n",
    "...or in general, for models where the embeddigns can have all different dimensions. In this library, these are: `TabMlp`, `TabResNet` and `TabNet`\n",
    "\n",
    "As shown in Figure, this is an encoder-encoder approach where we learn to predict values in the incoming data that have been masked. However, as I mentioned before, our implementation is a bit different, and the masking occurs in th embedding space. \n",
    "\n",
    "Nonetheless, the code below illustrates how to use this encoder-decoder approach with `pytorch-widedeep`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f63be691",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from pytorch_widedeep import Trainer\n",
    "from pytorch_widedeep.models import TabMlp, WideDeep\n",
    "from pytorch_widedeep.metrics import Accuracy\n",
    "from pytorch_widedeep.datasets import load_adult\n",
    "from pytorch_widedeep.preprocessing import TabPreprocessor\n",
    "from pytorch_widedeep.self_supervised_training import EncoderDecoderTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1aebb204",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = load_adult(as_frame=True)\n",
    "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
    "df.drop(\"income\", axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "16aab051",
   "metadata": {},
   "outputs": [],
   "source": [
    "# one could chose to use a validation set for early stopping, hyperparam\n",
    "# optimization, etc. This is just an example, so we simply use train/test\n",
    "# split\n",
    "df_tr, df_te = train_test_split(df, test_size=0.2, stratify=df.income_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9211ea45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>educational_num</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9042</th>\n",
       "      <td>26</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>250551</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Craft-repair</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25322</th>\n",
       "      <td>50</td>\n",
       "      <td>Private</td>\n",
       "      <td>34832</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Tech-support</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>15024</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       age  workclass  fnlwgt  education  educational_num      marital_status  \\\n",
       "9042    26  Local-gov  250551    HS-grad                9  Married-civ-spouse   \n",
       "25322   50    Private   34832  Bachelors               13  Married-civ-spouse   \n",
       "\n",
       "         occupation relationship   race gender  capital_gain  capital_loss  \\\n",
       "9042   Craft-repair    Own-child  Black   Male             0             0   \n",
       "25322  Tech-support      Husband  White   Male         15024             0   \n",
       "\n",
       "       hours_per_week native_country  income_label  \n",
       "9042               40  United-States             0  \n",
       "25322              40  United-States             1  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_tr.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dfee4efb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# As always, we need to define which cols will be represented as embeddings\n",
    "# and which one will be continuous features\n",
    "cat_embed_cols = [\n",
    "    \"workclass\",\n",
    "    \"education\",\n",
    "    \"marital_status\",\n",
    "    \"occupation\",\n",
    "    \"relationship\",\n",
    "    \"race\",\n",
    "    \"gender\",\n",
    "    \"capital_gain\",\n",
    "    \"capital_loss\",\n",
    "    \"native_country\",\n",
    "]\n",
    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
    "target_col = \"income_label\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "00e813bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/javierrodriguezzaurin/Projects/pytorch-widedeep/pytorch_widedeep/preprocessing/tab_preprocessor.py:358: UserWarning: Continuous columns will not be normalised\n",
      "  warnings.warn(\"Continuous columns will not be normalised\")\n"
     ]
    }
   ],
   "source": [
    "# We prepare the data to be passed to the model\n",
    "tab_preprocessor = TabPreprocessor(\n",
    "    cat_embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
    ")\n",
    "X_tab = tab_preprocessor.fit_transform(df_tr)\n",
    "target = df_tr[target_col].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "880be80c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1,  1,  1,  1,  1,  1,  1,  1,  1,  1, 26, 40],\n",
       "       [ 2,  2,  1,  2,  2,  2,  1,  2,  1,  1, 50, 40],\n",
       "       [ 2,  1,  1,  3,  2,  2,  1,  1,  2,  1, 39, 46],\n",
       "       [ 2,  3,  2,  4,  1,  2,  2,  1,  1,  1, 17, 10],\n",
       "       [ 3,  4,  2,  1,  1,  2,  1,  1,  1,  1, 32, 20]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_tab[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "442a7b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We define a model that will act as the encoder in the encoder/decoder\n",
    "# architecture. This could be any of: TabMlp, TabResnet or TabNet\n",
    "tab_mlp = TabMlp(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ee0eacb5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TabMlp(\n",
       "  (cat_embed): DiffSizeCatEmbeddings(\n",
       "    (embed_layers): ModuleDict(\n",
       "      (emb_layer_workclass): Embedding(10, 5, padding_idx=0)\n",
       "      (emb_layer_education): Embedding(17, 8, padding_idx=0)\n",
       "      (emb_layer_marital_status): Embedding(8, 5, padding_idx=0)\n",
       "      (emb_layer_occupation): Embedding(16, 7, padding_idx=0)\n",
       "      (emb_layer_relationship): Embedding(7, 4, padding_idx=0)\n",
       "      (emb_layer_race): Embedding(6, 4, padding_idx=0)\n",
       "      (emb_layer_gender): Embedding(3, 2, padding_idx=0)\n",
       "      (emb_layer_capital_gain): Embedding(124, 24, padding_idx=0)\n",
       "      (emb_layer_capital_loss): Embedding(98, 21, padding_idx=0)\n",
       "      (emb_layer_native_country): Embedding(42, 13, padding_idx=0)\n",
       "    )\n",
       "    (embedding_dropout): Dropout(p=0.0, inplace=False)\n",
       "  )\n",
       "  (cont_norm): Identity()\n",
       "  (encoder): MLP(\n",
       "    (mlp): Sequential(\n",
       "      (dense_layer_0): Sequential(\n",
       "        (0): Linear(in_features=95, out_features=200, bias=True)\n",
       "        (1): ReLU(inplace=True)\n",
       "        (2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (dense_layer_1): Sequential(\n",
       "        (0): Linear(in_features=200, out_features=100, bias=True)\n",
       "        (1): ReLU(inplace=True)\n",
       "        (2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tab_mlp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dd922f85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If we do not pass a custom decoder, which is perfectly possible via the\n",
    "# decoder param,  the EncoderDecoderTrainer will automatically build a\n",
    "# decoder which will be the 'mirror' image of the encoder\n",
    "encoder_decoder_trainer = EncoderDecoderTrainer(encoder=tab_mlp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "75311d4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EncoderDecoderModel(\n",
       "  (encoder): TabMlp(\n",
       "    (cat_embed): DiffSizeCatEmbeddings(\n",
       "      (embed_layers): ModuleDict(\n",
       "        (emb_layer_workclass): Embedding(10, 5, padding_idx=0)\n",
       "        (emb_layer_education): Embedding(17, 8, padding_idx=0)\n",
       "        (emb_layer_marital_status): Embedding(8, 5, padding_idx=0)\n",
       "        (emb_layer_occupation): Embedding(16, 7, padding_idx=0)\n",
       "        (emb_layer_relationship): Embedding(7, 4, padding_idx=0)\n",
       "        (emb_layer_race): Embedding(6, 4, padding_idx=0)\n",
       "        (emb_layer_gender): Embedding(3, 2, padding_idx=0)\n",
       "        (emb_layer_capital_gain): Embedding(124, 24, padding_idx=0)\n",
       "        (emb_layer_capital_loss): Embedding(98, 21, padding_idx=0)\n",
       "        (emb_layer_native_country): Embedding(42, 13, padding_idx=0)\n",
       "      )\n",
       "      (embedding_dropout): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "    (cont_norm): Identity()\n",
       "    (encoder): MLP(\n",
       "      (mlp): Sequential(\n",
       "        (dense_layer_0): Sequential(\n",
       "          (0): Linear(in_features=95, out_features=200, bias=True)\n",
       "          (1): ReLU(inplace=True)\n",
       "          (2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (dense_layer_1): Sequential(\n",
       "          (0): Linear(in_features=200, out_features=100, bias=True)\n",
       "          (1): ReLU(inplace=True)\n",
       "          (2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (decoder): TabMlpDecoder(\n",
       "    (decoder): MLP(\n",
       "      (mlp): Sequential(\n",
       "        (dense_layer_0): Sequential(\n",
       "          (0): Linear(in_features=100, out_features=200, bias=True)\n",
       "          (1): ReLU(inplace=True)\n",
       "          (2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (dense_layer_1): Sequential(\n",
       "          (0): Linear(in_features=200, out_features=95, bias=True)\n",
       "          (1): ReLU(inplace=True)\n",
       "          (2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (masker): RandomObfuscator()\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# let's have a look to the encoder_decoder_model (aka ed_model)\n",
    "encoder_decoder_trainer.ed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7442924a",
   "metadata": {},
   "source": [
    "Ignoring the `masker`, which just...well...masks, the `ed_model` consists of: \n",
    "1. An encoder model that is a `TabMlp` model that is in itself comprised by an Embedding layer (or rather a collection of them, referred as `cat_and_cont_embed`) and an encoder (a simple MLP, referred as `encoder`)\n",
    "2. A decoder which is just an \"inverted\" MLP (referred as `decoder`) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "864955bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 82.90it/s, loss=4.07]\n",
      "epoch 2: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 89.87it/s, loss=3.09]\n",
      "epoch 3: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 92.86it/s, loss=2.53]\n",
      "epoch 4: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 91.24it/s, loss=2.09]\n",
      "epoch 5: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 91.38it/s, loss=1.78]\n"
     ]
    }
   ],
   "source": [
    "# And we just...pretrain\n",
    "encoder_decoder_trainer.pretrain(X_tab, n_epochs=5, batch_size=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd803ee3",
   "metadata": {},
   "source": [
    "At this point we have two options, we could either save the model for later use or we could continue to supervised training. The latter is rather simple, after running:\n",
    "\n",
    "```python\n",
    "encoder_decoder_trainer.pretrain(X_tab, n_epochs=5, batch_size=256)\n",
    "```\n",
    "\n",
    "you just have to\n",
    "\n",
    "```python\n",
    "model = WideDeep(deeptabular=tab_mlp)\n",
    "trainer = Trainer(model=model, objective=\"binary\", metrics=[Accuracy])\n",
    "\n",
    "trainer.fit(X_tab=X_tab, target=target, n_epochs=5, batch_size=256)\n",
    "\n",
    "# And, you know...we get a test metric\n",
    "X_tab_te = tab_preprocessor.transform(df_te)\n",
    "target_te = df_te[target_col].values\n",
    "\n",
    "preds = trainer.predict(X_tab=X_tab_te)\n",
    "test_acc = accuracy_score(target_te, preds)\n",
    "```\n",
    "\n",
    "Let's say that in any case, we are 'decent' scientists/people and we want to save the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "793e6470",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_decoder_trainer.save(\n",
    "    path=\"pretrained_weights\", model_filename=\"encoder_decoder_model.pt\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a64db21a",
   "metadata": {},
   "source": [
    "some time has passed..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0212e8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_decoder_model = torch.load(\"pretrained_weights/encoder_decoder_model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b60a6ab",
   "metadata": {},
   "source": [
    "Now, **AND THIS IS IMPORTANT** We have loaded the encoder AND the decoder. To proceed to the supervised training we ONLY need the encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "28ff6876",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_encoder = encoder_decoder_model.encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5955bb15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TabMlp(\n",
       "  (cat_embed): DiffSizeCatEmbeddings(\n",
       "    (embed_layers): ModuleDict(\n",
       "      (emb_layer_workclass): Embedding(10, 5, padding_idx=0)\n",
       "      (emb_layer_education): Embedding(17, 8, padding_idx=0)\n",
       "      (emb_layer_marital_status): Embedding(8, 5, padding_idx=0)\n",
       "      (emb_layer_occupation): Embedding(16, 7, padding_idx=0)\n",
       "      (emb_layer_relationship): Embedding(7, 4, padding_idx=0)\n",
       "      (emb_layer_race): Embedding(6, 4, padding_idx=0)\n",
       "      (emb_layer_gender): Embedding(3, 2, padding_idx=0)\n",
       "      (emb_layer_capital_gain): Embedding(124, 24, padding_idx=0)\n",
       "      (emb_layer_capital_loss): Embedding(98, 21, padding_idx=0)\n",
       "      (emb_layer_native_country): Embedding(42, 13, padding_idx=0)\n",
       "    )\n",
       "    (embedding_dropout): Dropout(p=0.0, inplace=False)\n",
       "  )\n",
       "  (cont_norm): Identity()\n",
       "  (encoder): MLP(\n",
       "    (mlp): Sequential(\n",
       "      (dense_layer_0): Sequential(\n",
       "        (0): Linear(in_features=95, out_features=200, bias=True)\n",
       "        (1): ReLU(inplace=True)\n",
       "        (2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (dense_layer_1): Sequential(\n",
       "        (0): Linear(in_features=200, out_features=100, bias=True)\n",
       "        (1): ReLU(inplace=True)\n",
       "        (2): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretrained_encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "63e58e2e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|██████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 88.04it/s, loss=0.374, metrics={'acc': 0.8253}]\n",
      "epoch 2: 100%|██████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 85.63it/s, loss=0.324, metrics={'acc': 0.8491}]\n",
      "epoch 3: 100%|██████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 87.56it/s, loss=0.301, metrics={'acc': 0.8608}]\n",
      "epoch 4: 100%|███████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 73.38it/s, loss=0.29, metrics={'acc': 0.8655}]\n",
      "epoch 5: 100%|██████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 78.68it/s, loss=0.284, metrics={'acc': 0.8686}]\n",
      "predict: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 173.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8730678677449074\n"
     ]
    }
   ],
   "source": [
    "# and as always, ANY supervised model in this library has to go throuth the WideDeep class:\n",
    "model = WideDeep(deeptabular=pretrained_encoder)\n",
    "trainer = Trainer(model=model, objective=\"binary\", metrics=[Accuracy])\n",
    "\n",
    "trainer.fit(X_tab=X_tab, target=target, n_epochs=5, batch_size=256)\n",
    "\n",
    "X_tab_te = tab_preprocessor.transform(df_te)\n",
    "target_te = df_te[target_col].values\n",
    "\n",
    "preds = trainer.predict(X_tab=X_tab_te)\n",
    "test_acc = accuracy_score(target_te, preds)\n",
    "print(test_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69658ac7",
   "metadata": {},
   "source": [
    "As we mentioned before, we can also use a `TabResNet` or `TabNet` model and a custom decoder. Let's have a look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ff91ee50",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_widedeep.models import TabResnet as TabResnetEncoder, TabResnetDecoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "845bc1c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_encoder = TabResnetEncoder(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    continuous_cols=continuous_cols,\n",
    "    blocks_dims=[200, 100, 100],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26fe324",
   "metadata": {},
   "source": [
    "let's have a look to the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1d93e92a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TabResnet(\n",
       "  (cat_embed): DiffSizeCatEmbeddings(\n",
       "    (embed_layers): ModuleDict(\n",
       "      (emb_layer_workclass): Embedding(10, 5, padding_idx=0)\n",
       "      (emb_layer_education): Embedding(17, 8, padding_idx=0)\n",
       "      (emb_layer_marital_status): Embedding(8, 5, padding_idx=0)\n",
       "      (emb_layer_occupation): Embedding(16, 7, padding_idx=0)\n",
       "      (emb_layer_relationship): Embedding(7, 4, padding_idx=0)\n",
       "      (emb_layer_race): Embedding(6, 4, padding_idx=0)\n",
       "      (emb_layer_gender): Embedding(3, 2, padding_idx=0)\n",
       "      (emb_layer_capital_gain): Embedding(124, 24, padding_idx=0)\n",
       "      (emb_layer_capital_loss): Embedding(98, 21, padding_idx=0)\n",
       "      (emb_layer_native_country): Embedding(42, 13, padding_idx=0)\n",
       "    )\n",
       "    (embedding_dropout): Dropout(p=0.0, inplace=False)\n",
       "  )\n",
       "  (cont_norm): Identity()\n",
       "  (encoder): DenseResnet(\n",
       "    (dense_resnet): Sequential(\n",
       "      (lin_inp): Linear(in_features=95, out_features=200, bias=False)\n",
       "      (bn_inp): BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (block_0): BasicBlock(\n",
       "        (resize): Sequential(\n",
       "          (0): Linear(in_features=200, out_features=100, bias=False)\n",
       "          (1): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (lin1): Linear(in_features=200, out_features=100, bias=False)\n",
       "        (bn1): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (leaky_relu): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (dp): Dropout(p=0.1, inplace=False)\n",
       "        (lin2): Linear(in_features=100, out_features=100, bias=False)\n",
       "        (bn2): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (block_1): BasicBlock(\n",
       "        (lin1): Linear(in_features=100, out_features=100, bias=False)\n",
       "        (bn1): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (leaky_relu): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (dp): Dropout(p=0.1, inplace=False)\n",
       "        (lin2): Linear(in_features=100, out_features=100, bias=False)\n",
       "        (bn2): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resnet_encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad5eac07",
   "metadata": {},
   "source": [
    "As we can see, the tensor we are trying to reconstruct, the embeddings, is of size `94` (this number is stored in the attribute: `esnet_encoder.cat_and_cont_embed.output_dim`), with that information we could build or own decoder as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a2379a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for all possible params see the docs\n",
    "resnet_decoder = TabResnetDecoder(\n",
    "    embed_dim=resnet_encoder.cat_out_dim + resnet_encoder.cont_out_dim,\n",
    "    blocks_dims=[100, 100, 200],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "24a1f3cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TabResnetDecoder(\n",
       "  (decoder): DenseResnet(\n",
       "    (dense_resnet): Sequential(\n",
       "      (block_0): BasicBlock(\n",
       "        (lin1): Linear(in_features=100, out_features=100, bias=False)\n",
       "        (bn1): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (leaky_relu): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (dp): Dropout(p=0.1, inplace=False)\n",
       "        (lin2): Linear(in_features=100, out_features=100, bias=False)\n",
       "        (bn2): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (block_1): BasicBlock(\n",
       "        (resize): Sequential(\n",
       "          (0): Linear(in_features=100, out_features=200, bias=False)\n",
       "          (1): BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (lin1): Linear(in_features=100, out_features=200, bias=False)\n",
       "        (bn1): BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (leaky_relu): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (dp): Dropout(p=0.1, inplace=False)\n",
       "        (lin2): Linear(in_features=200, out_features=200, bias=False)\n",
       "        (bn2): BatchNorm1d(200, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (reconstruction_layer): Linear(in_features=200, out_features=95, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resnet_decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "877a5414",
   "metadata": {},
   "source": [
    "and now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9607b5e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:03<00:00, 46.89it/s, loss=1.52]\n",
      "epoch 2: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:03<00:00, 46.78it/s, loss=0.81]\n",
      "epoch 3: 100%|████████████████████████████████████████████████████████████████████████████████████| 153/153 [00:03<00:00, 39.82it/s, loss=0.56]\n",
      "epoch 4: 100%|███████████████████████████████████████████████████████████████████████████████████| 153/153 [00:03<00:00, 46.73it/s, loss=0.417]\n",
      "epoch 5: 100%|███████████████████████████████████████████████████████████████████████████████████| 153/153 [00:03<00:00, 46.24it/s, loss=0.329]\n"
     ]
    }
   ],
   "source": [
    "ec_trainer = EncoderDecoderTrainer(\n",
    "    encoder=resnet_encoder,\n",
    "    decoder=resnet_decoder,\n",
    "    masked_prob=0.2,\n",
    ")\n",
    "ec_trainer.pretrain(X_tab, n_epochs=5, batch_size=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5401740f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|██████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 58.63it/s, loss=0.335, metrics={'acc': 0.8442}]\n",
      "epoch 2: 100%|███████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 58.02it/s, loss=0.296, metrics={'acc': 0.864}]\n",
      "epoch 3: 100%|██████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 55.91it/s, loss=0.283, metrics={'acc': 0.8687}]\n",
      "epoch 4: 100%|███████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 55.00it/s, loss=0.276, metrics={'acc': 0.871}]\n",
      "epoch 5: 100%|██████████████████████████████████████████████████████████| 153/153 [00:02<00:00, 51.95it/s, loss=0.272, metrics={'acc': 0.8732}]\n",
      "predict: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 120.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8725560446309756\n"
     ]
    }
   ],
   "source": [
    "# and as always, ANY supervised model in this library has to go throuth the WideDeep class:\n",
    "model = WideDeep(deeptabular=resnet_encoder)\n",
    "trainer = Trainer(model=model, objective=\"binary\", metrics=[Accuracy])\n",
    "\n",
    "trainer.fit(X_tab=X_tab, target=target, n_epochs=5, batch_size=256)\n",
    "\n",
    "X_tab_te = tab_preprocessor.transform(df_te)\n",
    "target_te = df_te[target_col].values\n",
    "\n",
    "preds = trainer.predict(X_tab=X_tab_te)\n",
    "test_acc = accuracy_score(target_te, preds)\n",
    "print(test_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "729606ac-d509-42d6-894c-8d69d2c12bbc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
